# coding: utf-8
"""
网络结构中一些常用的block的集合
"""
import torch.nn as nn
import torch.nn.functional as F

class DBL(nn.Module):
    def __init__(self, ch_in, ch_out, kernel=3, pad=1, stride=1, prelu=True):
        super(DBL, self).__init__()

        self.conv = nn.Conv2d(ch_in, ch_out, kernel_size=kernel, padding=pad, stride=stride)
        self.bn = nn.BatchNorm2d(ch_out)
        self.relu = None
        if prelu:
            self.relu = nn.PReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)

        if self.relu:
            x = self.relu(x)
        return x

class dark_res(nn.Module):
    def __init__(self, ch_in):
        super(dark_res, self).__init__()
        self.conv1 = DBL(ch_in, ch_in // 2, kernel=1, stride=1, pad=0)
        self.conv2 = DBL(ch_in // 2, ch_in, kernel=3, stride=1, pad=1)

    def forward(self, x):
        y = self.conv1(x)
        y = self.conv2(y)

        return x + y

